import json
from typing import Dict, Any, List
from langchain.agents import AgentType, initialize_agent, Tool
from langchain.chat_models import ChatOpenAI
from langchain.schema import SystemMessage
from langchain.prompts import MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
# import openai
from langchain_ollama import OllamaLLM

from transformers import T5ForConditionalGeneration, T5Tokenizer
TRANSFORMERS_AVAILABLE = True
import warnings
warnings.filterwarnings("ignore")

# 1. 定义 Agent 的工具函数

def read_issue_content(issue_data: dict) -> str:
    """
    从模拟的 Issue 数据中提取标题和描述正文
    """
    try:
        title = issue_data.get('title', '')
        body = issue_data.get('body', '')
        return f"标题: {title}\n描述: {body}"
    except Exception as e:
        return f"读取Issue内容时出错: {str(e)}"


class T5IssueClassifier:
    """使用 T5 模型进行 Issue 分类"""

    def __init__(self, model_name: str = "t5-small"):
        """初始化 T5 模型和分词器"""
        self.tokenizer = T5Tokenizer.from_pretrained('../T5_small')
        self.model = T5ForConditionalGeneration.from_pretrained('../T5_small')

    def classify_with_t5(self, text_content: str) -> str:
        """
        使用 T5 模型对 Issue 内容进行分类
        """

        # try:
        # 构建分类提示
        prompt = f"""你是一个专业的问题分类助手。请将用户提供的问题进行分类
        分为以下三种类型之一：
        Bug（系统错误、故障、异常行为等）
        Feature Request（新功能请求、功能改进建议等）
        Question（使用问题、咨询、帮助请求等)
        只输出分类的结果Bug，Feature Request，Question，不允许输出额外的信息
        例如：创建进程时出现错误
        输出：Bug
        根据以上示例对用户提问内容:{text_content}进行分类输出。
        """

        # 编码输入
        inputs = self.tokenizer.encode(prompt, return_tensors="pt", max_length=128, truncation=True)

        # 生成输出
        outputs = self.model.generate(
            inputs,
            max_length=10,
            num_beams=2,
            early_stopping=True
        )

        # 解码输出
        decoded_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        # print(decoded_output)
        # 映射到类别
        return decoded_output

    @staticmethod
    def map_t5_output_to_category(t5_output: str) -> str:
        """将 T5 输出映射到我们的类别系统"""
        output_lower = t5_output.lower()

        if any(keyword in output_lower for keyword in ['bug', 'error', 'issue']):
            return "Bug"
        elif any(keyword in output_lower for keyword in ['feature', 'request', 'enhancement']):
            return "Feature Request"
        elif any(keyword in output_lower for keyword in ['question', 'help', 'support']):
            return "Question"



# 创建 T5 分类器实例
t5_classifier = T5IssueClassifier()
def categorize_issue(text_content: str) -> str:
    """
    使用 T5 模型将 Issue 内容分类
    """
    t5_result = t5_classifier.classify_with_t5(text_content)
    output = t5_classifier.map_t5_output_to_category(t5_result)
    return output


def assign_developer(category: str) -> str:
    """
    根据分类结果分配负责人
    """
    developer_mapping = {
        "Bug": "张三",
        "Feature Request": "李四",
        "Question": "王五"
    }

    return developer_mapping.get(category)


# 2. 构建 LangChain Agent 执行逻辑

class IssueProcessingAgent:
    def __init__(self):
        """
        初始化Agent
        """
        # 创建工具列表
        self.tools = [
            Tool(
                name="read_issue_content",
                func=read_issue_content,
                description="从Issue数据中提取标题和描述正文"
            ),
            Tool(
                name="categorize_issue",
                func=categorize_issue,
                description="使用T5模型将Issue内容分类为 Bug, Feature Request, Question 中的一种"
            ),
            Tool(
                name="assign_developer",
                func=assign_developer,
                description="根据分类结果分配负责人, Bug:张三, Feature Request:李四, Question:王五"
            )
        ]

        # 初始化LLM  使用ollamallm调用本地模型
        self.llm = OllamaLLM(base_url="127.0.0.1:11434", model="deepseek-r1:7b")
        # 系统提示词
        system_message = SystemMessage(
            content="""你是一个专业的Issue处理助手。你的任务是：
            1. 读取和理解Issue的内容
            2. 使用T5模型准确分类Issue的类型
            3. 分配合适的开发者进行处理

            请按照步骤思考并调用合适的工具来完成工作。"""
        )

        # 创建记忆
        # self.memory = ConversationBufferMemory(
        #     memory_key="chat_history",
        #     return_messages=True
        # )

        # 初始化Agent
        self.agent = initialize_agent(
            tools=self.tools,
            llm=self.llm,
            agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
            # memory=self.memory,
            # system_message=system_message,
            # agent_kwargs={
            #     "extra_prompt_messages": [MessagesPlaceholder(variable_name="chat_history")]
            # }
        )

    def process_issue(self, issue_data: Dict[str, Any]):
        """
        处理Issue的主要方法
        """
        # 构建任务描述
        # print(issue_data)
        task = f"""
        请处理以下Issue数据格式为json：
         {json.dumps(issue_data, ensure_ascii=False, indent=2)}

        请调用自定义工具执行以下步骤：
        1. 读取Issue内容
        2. 使用T5模型分类Issue类型
        3. 根据分类结果分配负责人

        仅输出最终格式：
        - 分类结果: [分类]
        - 分配负责人: [负责人]
        """

        # 执行Agent
        result = self.agent.run(task)
        print(result)


def create_sample_issues() -> List[Dict[str, Any]]:
    """创建模拟的Issue数据"""
    return [
        {
            "title": "系统登录时出现500错误",
            "body": "用户尝试登录时系统返回500内部服务器错误，需要立即修复。",
        },
        {
            "title": "希望增加黑暗模式支持",
            "body": "建议为应用添加黑暗模式功能，提升夜间使用体验。",
        },
        {
            "title": "如何重置密码？",
            "body": "请问忘记密码后应该如何重置？在设置中没有找到相关选项。",
        }
    ]


def main():
    """主函数演示Agent的使用"""

    # 创建Agent实例

    # 获取模拟数据
    sample_issues = create_sample_issues()

    print("开始处理Issue...")
    print("=" * 50)

    # 测试 T5 分类器
    print("测试 T5 分类器:")
    test_text = "系统登录时出现错误"
    category = categorize_issue(test_text)
    print(f"测试文本: '{test_text}' -> 分类: {category}")
    print("-" * 30)

    # Agent处理每个Issue
    issue_agent = IssueProcessingAgent()
    for i, issue in enumerate(sample_issues, 1):
        print(f"\n处理第 {i} 个Issue:")
        print(f"标题: {issue['title']}")
        issue_agent.process_issue(issue)
        print("-" * 50)

if __name__ == '__main__':
    main()